package learner.hmm;

import java.util.Map;

public class AlgorithmValueCalculator {

	private Alpha alpha;

	private Beta beta;

	private Matrix<Integer, Integer> a;

	private Matrix<Integer, MovementVector> b;

	private Movement movement;

	public AlgorithmValueCalculator(Matrix<Integer, Integer> a,
			Matrix<Integer, MovementVector> b, Map<Integer, Double> pi,
			Movement movement) {
		alpha = new Alpha(a, b, pi, movement);
		beta = new Beta(a, b, movement);
		this.a = a;
		this.b = b;
		this.movement = movement;
	}

	public AlgorithmValueCalculator(HMMModel model, Movement m) {
		this(new Matrix<Integer, Integer>(model.getA()),
				new Matrix<Integer, MovementVector>(model.getB()), model
						.getPi(), m);
	}

	/**
	 * zie lecture slide : baum-welch algorithm : expectation step(1)
	 * 
	 * @param k
	 *            The time for which you want to know the psi
	 * @param i
	 *            The state corresponding to time k
	 * @param j
	 *            The state corresponding to time k + 1
	 * @return
	 */
	public double getPsiForTime(int k, int i, int j) {
		double numerator = alpha.get(k, i) * a.get(i, j)
				* b.get(j, movement.getMovementForTime(k + 1))
				* beta.get(k + 1, j);
		double denominator = 0.0;
		for (int stateI = 0; stateI < BaumWelchLearner.NB_OF_STATES; stateI++) {
			for (int stateJ = 0; stateJ < BaumWelchLearner.NB_OF_STATES; stateJ++) {
				denominator += alpha.get(k, stateI) * a.get(stateI, stateJ)
						* b.get(stateJ, movement.getMovementForTime(k + 1))
						* beta.get(k + 1, stateJ);
			}
		}
		return numerator / denominator;
	}

	/**
	 * zie lecture slide : baum-welch algorithm : expectation step(2)
	 * 
	 * @param k
	 *            The time for which you want to know the gamma
	 * @param i
	 *            The state for which you want to know gamma.
	 * @return
	 */
	public double getGammaForTime(int k, int i) {
		double gamma_k_i = 0;
		if (k == 1) {
			for (int j = 0; j < BaumWelchLearner.NB_OF_STATES; j++) {
				gamma_k_i += getPsiForTime(k, i, j);
			}
		} else {
			for (int j = 0; j < BaumWelchLearner.NB_OF_STATES; j++) {
				gamma_k_i += getPsiForTime(k - 1, j, i);
			}
		}
		return gamma_k_i;
	}

	public double getChanceOnPattern() {
		double chance = 0.0;
		for (int i = 0; i < BaumWelchLearner.NB_OF_STATES; i++) {
			chance += alpha.get(movement.getSize(), i);
		}
		return chance;
	}
}
